# -*- coding: utf-8 -*-
# Functions for encryption and signing based on PKCS#1 v1.5
# See http://www.di-mgt.com.au/rsa_alg.html#pkcs1schemes
"""Functions for encryption and signing based on PKCS#1 v1.5.
See http://www.di-mgt.com.au/rsa_alg.html#pkcs1schemes

NOTICE: This program is used for learning, there may have bugs
and potential security issues. Please DO NOT use this program
in production environment.

WARNING: Do not pass any Python traceback information to users
for it may leaks the type and the position of a failure when
it occurred.
"""
from rsa.assertfunc import param_assert_str, param_assert_private_key, param_assert_public_key

__author__ = 'Kainan Zhu'
__email__ = 'dlangu0393@gmail.com'
__copyright__ = 'Copyright 2013, Kainan Zhu'

import os
import hashlib
from rsa.util import byte_length, bytes_to_int, int_to_bytes
from rsa.core import encrypt_int, decrypt_int
from rsa.common import DecryptionException, VerificationException, CryptoException


HASH_METHODS = {
    'MD5': hashlib.md5,
    'SHA-1': hashlib.sha1,
    'SHA-256': hashlib.sha256,
    'SHA-384': hashlib.sha384,
    'SHA-512': hashlib.sha512,
}

# ASN.1 codes that describe the hash algorithm used.
# See http://www.ietf.org/rfc/rfc3447.txt  [Page 42]
HASH_ASN1 = {
    'MD5': '\x30\x20\x30\x0c\x06\x08\x2a\x86\x48\x86\xf7\x0d\x02\x05\x05\x00\x04\x10',
    'SHA-1': '\x30\x21\x30\x09\x06\x05\x2b\x0e\x03\x02\x1a\x05\x00\x04\x14',
    'SHA-256': '\x30\x31\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\x01\x05\x00\x04\x20',
    'SHA-384': '\x30\x41\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\x02\x05\x00\x04\x30',
    'SHA-512': '\x30\x51\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\x03\x05\x00\x04\x40',
}


def encryption_padding(message, target_length):
    """
    Padding message for encryption
    :param message:
    :param target_length:
    :return:
        00 02 PADDING 00 MESSAGE
            in which PADDING is a string of TARGET_LENGTH-|MESSAGE|-3
            non-zero randomly-generated bytes(at least eight random bytes).
    """
    # Insure that there are enough space for constant bytes and at least eight random bytes
    if len(message) > target_length - 11:
        raise ValueError('target_length is too small to place the message')

    # Generate the random padding string
    #   which should not contain 0-byte.
    padding = ''
    padding_length = target_length - len(message) - 3
    while len(padding) < padding_length:
        length_needed = padding_length - len(padding)
        # Generate more and remove the 0-byte, then trim it
        padding += os.urandom(length_needed + 8).replace('\x00', '')[:length_needed]

    return ''.join(['\x00\x02', padding, '\x00', message])


def signing_padding(message, target_length):
    """
    Padding message for signing
    :param message:
    :param target_length:
    :return:
        00 01 PADDING 00 MESSAGE
        in which PADDING is a string of TARGET_LENGTH-|MESSAGE|-3 FF bytes
            (at least eight random bytes).
    """
    # Insure that there are enough space for constant bytes and at least eight random bytes
    if len(message) > target_length - 11:
        raise ValueError('target_length is too small to place the message')

    padding_length = target_length - len(message) - 3
    padding = '\xff' * padding_length

    return ''.join(['\x00\x01', padding, '\x00', message])


def encrypt(message, key):
    """
    Encrypt the message using key based on PKCS#1 v1.5
    :param message:
        The message to be encrypted.
    :param key:
        The public key used in encrypt.
    :return:
        A string of encrypted bytes.

    >>> from rsa import key
    >>> (pub_key, pri_key) = key.new_keys(256)
    >>> message = 'Hello'
    >>> cypher = encrypt(message, pub_key)

    """
    param_assert_str(message, 'message')
    param_assert_public_key(key, 'key')

    key_length = byte_length(key.n)
    padded_message = encryption_padding(message, key_length)
    int_message = bytes_to_int(padded_message)
    encrypted = encrypt_int(int_message, key.e, key.n)
    cypher = int_to_bytes(encrypted, key_length)

    return cypher


def decrypt(cypher, key):
    """
    Decrypt the cypher using key based on PKCS#1 v1.5
    :param cypher:
        The cypher to be decrypted.
    :param key:
        The private key used in decrypt.
    :return:
        A string of decrypted bytes.

    >>> from rsa import key
    >>> (pub_key, pri_key) = key.new_keys(256)
    >>> message = 'Hello'
    >>> cypher = encrypt(message, pub_key)
    >>> decrypt(cypher, pri_key)
    'Hello'

    """
    param_assert_str(cypher, 'cypher')
    param_assert_private_key(key, 'key')

    key_length = byte_length(key.n)
    int_cypher = bytes_to_int(cypher)
    decrypted = decrypt_int(int_cypher, key.d, key.n)
    plain = int_to_bytes(decrypted, key_length)

    if not plain[0:2] == '\x00\x02':
        raise DecryptionException('Decryption failed!')

    try:
        idx = plain.index('\x00', 2)
    except ValueError:
        raise DecryptionException('Decryption failed!')

    return plain[idx + 1:]


def sign(message, key, hash_method):
    """
    Sign the message with key using the hash method
    :param message:
        The message to be signed.
    :param key:
        The private key used in signing.
    :param hash_method:
        The hash method.
        Valid: 'MD5', 'SHA-1', 'SHA-256', 'SHA-384' or 'SHA-512'.
    :return:
        A string of message signature block.
    """
    param_assert_str(message, 'message')
    param_assert_private_key(key, 'key')
    param_assert_str(hash_method, 'hash_method')

    hash_method = hash_method.upper()
    if not hash_method in HASH_METHODS.keys():
        raise ValueError("Param hash_method should be one of 'MD5', 'SHA-1', 'SHA-256', 'SHA-384' or 'SHA-512'")

    # Calculate hash
    hash_instance = HASH_METHODS[hash_method]()
    hash_instance.update(message)
    hash_code = hash_instance.digest()

    # Generate sign plain text
    plain = HASH_ASN1[hash_method] + hash_code

    key_length = byte_length(key.n)
    padded_plain = signing_padding(plain, key_length)
    int_plain = bytes_to_int(padded_plain)
    encrypted = encrypt_int(int_plain, key.d, key.n)
    signature = int_to_bytes(encrypted, key_length)

    return signature


def verify(message, signature, key):
    """
    Verify if the signature matches the message based on PKCS#1 v1.5
    :param message:
        The signed message.
    :param signature:
        The signature block.
    :param key:
        The key used in verification.
    :return:
        A bool value.
    """
    param_assert_str(message, 'message')
    param_assert_str(signature, 'signature')
    param_assert_public_key(key, 'key')

    key_length = byte_length(key.n)
    int_signature = bytes_to_int(signature)
    decrypted = decrypt_int(int_signature, key.e, key.n)
    plain = int_to_bytes(decrypted, key_length)

    if not plain[0:2] == '\x00\x01':
        raise VerificationException('Verification failed!')

    try:
        idx = plain.index('\x00', 2)
    except ValueError:
        raise VerificationException('Verification failed!')

    # De-pad the plain
    plain = plain[idx + 1:]

    # Find hash method and hash
    hash_method = ''
    hash_code = ''
    for (method, asn1) in HASH_ASN1.items():
        if plain.startswith(asn1):
            hash_method = method
            hash_code = plain[len(asn1):]
            break

    if not hash_method in HASH_METHODS.keys():
        raise VerificationException('Verification failed!')

    # Calculate new hash
    hash_instance = HASH_METHODS[hash_method]()
    hash_instance.update(message)
    hash_code_now = hash_instance.digest()

    # Compare two hashes
    if not hash_code == hash_code_now:
        raise VerificationException('Verification failed!')

    return True


__all__ = ['encrypt', 'decrypt', 'sign', 'verify',
           'DecryptionException', 'VerificationException', 'CryptoException']
